import akshare as ak
import pandas as pd
import pickle
import os
from datetime import datetime, timedelta


class TradingCalendar:
    def __init__(self, cache_file=None, refresh_days=30):
        # 使用与模块相关的绝对路径作为默认缓存路径
        if cache_file is None:
            module_dir = os.path.dirname(os.path.abspath(__file__))
            self.cache_file = os.path.join(module_dir, 'trading_calendar.pkl')
        else:
            self.cache_file = cache_file
        
        self.refresh_days = refresh_days
        self.trading_days = None
        self.load_calendar()

    def load_calendar(self):
        """加载交易日历，优先使用缓存"""
        # 检查缓存是否存在且未过期
        if self._is_cache_valid():
            try:
                self._load_from_cache()
                print(f"从缓存加载交易日历: {self.cache_file}")
            except Exception as e:
                print(f"从缓存加载失败，将重新下载: {e}")
                self._download_and_cache()
        else:
            self._download_and_cache()
            print(f"重新下载并缓存交易日历到: {self.cache_file}")

    def _is_cache_valid(self):
        """检查缓存是否有效"""
        if not os.path.exists(self.cache_file):
            return False

        # 检查缓存文件的修改时间
        cache_time = datetime.fromtimestamp(os.path.getmtime(self.cache_file))
        if (datetime.now() - cache_time).days > self.refresh_days:
            return False

        return True

    def _load_from_cache(self):
        """从缓存加载数据"""
        try:
            with open(self.cache_file, 'rb') as f:
                self.trading_days = pickle.load(f)
        except (pickle.UnpicklingError, IOError, EOFError) as e:
            print(f"缓存文件读取错误: {e}")
            raise

    def _download_and_cache(self):
        """下载数据并缓存"""
        try:
            self.trading_days = ak.tool_trade_date_hist_sina()
            # 确保目录存在
            os.makedirs(os.path.dirname(self.cache_file), exist_ok=True)
            # 缓存到文件
            with open(self.cache_file, 'wb') as f:
                pickle.dump(self.trading_days, f)
        except Exception as e:
            print(f"下载交易日历失败: {e}")
            self.trading_days = None

    def is_trading_day(self, date_obj):
        """判断是否为交易日"""
        if self.trading_days is None:
            return False

        try:
            # 统一日期格式
            if isinstance(date_obj, str):
                # 验证日期字符串格式
                try:
                    trade_date = pd.to_datetime(date_obj).strftime('%Y%m%d')
                except ValueError:
                    print(f"无效的日期字符串格式: {date_obj}")
                    return False
            elif isinstance(date_obj, (datetime, pd.Timestamp)):
                trade_date = date_obj.strftime('%Y%m%d')
            elif isinstance(date_obj, datetime.date):
                trade_date = date_obj.strftime('%Y%m%d')
            else:
                print(f"不支持的日期类型: {type(date_obj)}")
                return False

            # 确保与交易日历中的数据类型一致
            if isinstance(self.trading_days['trade_date'].iloc[0], str):
                return trade_date in self.trading_days['trade_date'].values
            else:
                # 如果交易日历中的数据是日期类型，将trade_date转换为日期类型
                trade_date_dt = datetime.strptime(trade_date, '%Y%m%d').date()
                return any(date == trade_date_dt for date in self.trading_days['trade_date'].values)
        except Exception as e:
            print(f"日期格式错误: {e}")
            return False

    def get_trading_days_in_range(self, start_date, end_date):
        """获取指定范围内的所有交易日"""
        if self.trading_days is None:
            return []

        try:
            # 转换开始和结束日期为datetime对象
            start_dt = pd.to_datetime(start_date)
            end_dt = pd.to_datetime(end_date)
            
            # 检查交易日历中的数据类型并进行相应的比较
            if isinstance(self.trading_days['trade_date'].iloc[0], str):
                # 如果是字符串类型，转换为相同格式的字符串进行比较
                start_str = start_dt.strftime('%Y%m%d')
                end_str = end_dt.strftime('%Y%m%d')
                mask = (self.trading_days['trade_date'] >= start_str) & (self.trading_days['trade_date'] <= end_str)
            else:
                # 如果是日期类型，转换为date对象进行比较
                start_date_obj = start_dt.date()
                end_date_obj = end_dt.date()
                mask = (self.trading_days['trade_date'] >= start_date_obj) & (self.trading_days['trade_date'] <= end_date_obj)
            
            return self.trading_days[mask]['trade_date'].tolist()
        except Exception as e:
            print(f"查询失败: {e}")
            return []
            
    def get_next_trading_day(self, date_obj):
        """获取指定日期之后的下一个交易日"""
        if self.trading_days is None:
            return None
            
        try:
            # 统一日期格式
            if isinstance(date_obj, str):
                date = pd.to_datetime(date_obj)
            elif isinstance(date_obj, datetime):
                date = pd.to_datetime(date_obj)
            elif isinstance(date_obj, datetime.date):
                date = pd.to_datetime(date_obj)
            else:
                return None
                
            # 检查交易日历中的数据类型并进行相应的比较
            if isinstance(self.trading_days['trade_date'].iloc[0], str):
                # 如果是字符串类型，转换为相同格式的字符串进行比较
                date_str = date.strftime('%Y%m%d')
                mask = self.trading_days['trade_date'] > date_str
            else:
                # 如果是日期类型，转换为date对象进行比较
                date_obj = date.date()
                mask = self.trading_days['trade_date'] > date_obj
                
            future_days = self.trading_days[mask]['trade_date'].tolist()
            
            if future_days:
                return future_days[0]  # 返回第一个大于指定日期的交易日
            return None
        except Exception as e:
            print(f"获取下一个交易日失败: {e}")
            return None
            
    def get_previous_trading_day(self, date_obj):
        """获取指定日期之前的上一个交易日"""
        if self.trading_days is None:
            return None
            
        try:
            # 统一日期格式
            if isinstance(date_obj, str):
                date = pd.to_datetime(date_obj)
            elif isinstance(date_obj, datetime):
                date = pd.to_datetime(date_obj)
            elif isinstance(date_obj, datetime.date):
                date = pd.to_datetime(date_obj)
            else:
                return None
                
            # 检查交易日历中的数据类型并进行相应的比较
            if isinstance(self.trading_days['trade_date'].iloc[0], str):
                # 如果是字符串类型，转换为相同格式的字符串进行比较
                date_str = date.strftime('%Y%m%d')
                mask = self.trading_days['trade_date'] < date_str
            else:
                # 如果是日期类型，转换为date对象进行比较
                date_obj = date.date()
                mask = self.trading_days['trade_date'] < date_obj
                
            past_days = self.trading_days[mask]['trade_date'].tolist()
            
            if past_days:
                return past_days[-1]  # 返回最后一个小于指定日期的交易日
            return None
        except Exception as e:
            print(f"获取上一个交易日失败: {e}")
            return None

    def refresh_calendar(self):
        """强制刷新交易日历"""
        self._download_and_cache()


# 使用示例
if __name__ == "__main__":
    # 创建交易日历实例（会自动缓存）
    # calendar = TradingCalendar()

    # # 使用当前年份或过去年份的日期进行测试
    # current_year = "2025"
    # test_dates = [f"{current_year}-10-01", f"{current_year}-10-08", f"{current_year}-10-09"]
    # for date in test_dates:
    #     result = calendar.is_trading_day(date)
    #     print(f"{date}: {'交易日' if result else '非交易日'}")

    # print("\n" + "=" * 50)

    # # 测试日期范围
    # start_date = f"{current_year}-10-01"
    # end_date = f"{current_year}-10-10"
    # trading_days = calendar.get_trading_days_in_range(start_date, end_date)
    # print(f"{start_date} 到 {end_date} 的交易日:")
    # print(trading_days)
    pass